rm(list=ls())
library(fields); library(maps); library(ncdf);library(abind)
library(geosphere);library(mapdata)
setwd('/Volumes/Disk2/JJA_ozone_MAM/Harvard_Dataverse')

#----- load functions -----------------------
source("./Function/get_geo.R")
source("./Function/get_met.R")
source('./Function/read_detrend.R')
source('./Function/eof.mca.R')
source('./Function/cov4gappy.R')
source('./Function/Henderson.R')
source('./Function/filled.contour3.R')

name="./Data/land.nc"
datafile = open.ncdf(name)
lon.land = get.var.ncdf(datafile, varid='lon')-360
lat.land = get.var.ncdf(datafile, varid='lat')
lat.land=rev(lat.land)
land = get.var.ncdf(datafile, varid='land')
land=land[,length(lat.land):1]
close.ncdf(datafile)


################read SST ##########################
for(year in 1979:2008){
	print(year)
	name=paste('./Data/SST/360x180/amipobs_sst_360x180_',year,'.nc',sep='')
	datafile = open.ncdf(name)
	lon1=get.var.ncdf(datafile,varid='longitude')
	lat1=get.var.ncdf(datafile,varid='latitude')
	time=get.var.ncdf(datafile,varid='time')
	ind1=which(lon1>=180 & lon1<=360)
	ind2=which(lat1>=0 & lat1<=90)
	data=get.var.ncdf(datafile,start=c(ind1[1],ind2[1],1),count=c(length(ind1),length(ind2),length(time)),varid='tos')
	close.ncdf(datafile)
	lon1=lon1[ind1]
	lat1=lat1[ind2]
	if(year==1979) {
		obs_SST=data
	} else {
		obs_SST=abind(obs_SST,data,along=3)
	}
}

ss=seq(as.Date('1979-01-01'),as.Date('2008-12-01'),by='month')
date=cbind(as.numeric(format(ss,"%Y")),as.numeric(format(ss,"%m")))
MAM_SST=array(NA,c(dim(obs_SST)[1],dim(obs_SST)[2],30))
JJA_SST=array(NA,c(dim(obs_SST)[1],dim(obs_SST)[2],30))
for (year in 1979:2008){
	ind=(date[,1]==year & date[,2]>=(4-1) & date[,2]<=(4+1))
	MAM_SST[,,year-1978]=apply(obs_SST[,,ind],c(1,2),mean,na.rm=TRUE)	
	ind=(date[,1]==year & date[,2]>=(7-1) & date[,2]<=(7+1))
	JJA_SST[,,year-1978]=apply(obs_SST[,,ind],c(1,2),mean,na.rm=TRUE)	
}

for(ii in 1:dim(MAM_SST)[1]){
	for(jj in 1:dim(MAM_SST)[2]){
		MAM_SST[ii,jj,]=mov.detrend(MAM_SST[ii,jj,],len=7)
		JJA_SST[ii,jj,]=mov.detrend(JJA_SST[ii,jj,],len=7)		
	}
}

lon.frame=lon1
lat.frame=lat1


#============================================================
workdir='./Data/AMIP/'
model.names=list.dirs(path=workdir,full.names=FALSE,recursive=FALSE)
for(kk in 1:length(model.names)){
print('=====================')
print(kk)
setwd(paste(workdir,model.names[kk],sep=''))
#(1) Read tas ---------------
files=(Sys.glob("tas_*.nc"))
files=sort(files)
for(ii in 1:length(files)){
	name=files[ii]
	print(name)
	scenarios=unlist(strsplit(name,"[._]"))
	ap=scenarios[6]
	ng=(unlist(strsplit(ap,"-")))
	time1=as.Date(paste(substr(ng[1],1,4),'-',substr(ng[1],5,6),'-01',sep=''))
	time2=as.Date(paste(substr(ng[2],1,4),'-',substr(ng[2],5,6),'-01',sep=''))
	time3=seq(time1,time2,by='month')
	date=cbind(as.numeric(format(time3,"%Y")),as.numeric(format(time3,"%m")))
	datafile = open.ncdf(name)
	lon=get.var.ncdf(datafile,varid='lon')
	lat=get.var.ncdf(datafile,varid='lat')
	ind1=which(lon>=180 & lon<=360)
	ind2=which(lat>=0 & lat<=90)
	time=get.var.ncdf(datafile,varid='time') 
	data=get.var.ncdf(datafile,start=c(ind1[1],ind2[1],1),count=c(length(ind1),length(ind2),length(time)),varid='tas')
	close.ncdf(datafile)
	lon.out=lon[ind1]-360
	lat.out=lat[ind2]
	if(ii==1) {
		amip_tas=data
		tas_time =date
	} else {
		amip_tas=abind(amip_tas,data,along=3)
		tas_time=rbind(tas_time,date)
	}
}

#(2) Read slp ---------------
files=(Sys.glob("psl_*.nc"))
files=sort(files)
for(ii in 1:length(files)){
	name=files[ii]
	print(name)
	scenarios=unlist(strsplit(name,"[._]"))
	ap=scenarios[6]
	ng=(unlist(strsplit(ap,"-")))
	time1=as.Date(paste(substr(ng[1],1,4),'-',substr(ng[1],5,6),'-01',sep=''))
	time2=as.Date(paste(substr(ng[2],1,4),'-',substr(ng[2],5,6),'-01',sep=''))
	time3=seq(time1,time2,by='month')
	date=cbind(as.numeric(format(time3,"%Y")),as.numeric(format(time3,"%m")))
	datafile = open.ncdf(name)
	lon=get.var.ncdf(datafile,varid='lon')
	lat=get.var.ncdf(datafile,varid='lat')
	ind1=which(lon>=180 & lon<=360)
	ind2=which(lat>=0 & lat<=90)
	time=get.var.ncdf(datafile,varid='time') 
	data=get.var.ncdf(datafile,start=c(ind1[1],ind2[1],1),count=c(length(ind1),length(ind2),length(time)),varid='psl')
	close.ncdf(datafile)
	lon.out=lon[ind1]-360
	lat.out=lat[ind2]
	if(ii==1) {
		amip_psl=data
		slp_time=date
	} else {
		amip_psl=abind(amip_psl,data,along=3)
		slp_time =rbind(slp_time,date)
	}
}


#(3) Define a function ---------------
read.amip=function(spdata,date,lon.out,lat.out,start_month,end_month){
	for(year in 1979:2008){
		ind=(date[,1]==year & date[,2]>=start_month & date[,2]<=end_month)
		ap=apply(spdata[,,ind],c(1,2),mean)
		if(year==1979){
			result=ap
		} else {
			result=abind(result,ap,along=3)
		}
	}
	return(result)	
}

#(4) temperature in the eastern United States ---------------
ind1=(lon.out>=-100 & lon.out <=-65)
ind2=(lat.out>=31 & lat.out <=50)
mask=sp.dissolve(land,lon.land,lat.land,lon.out,lat.out)
mask[mask<=0.25]=0;space.ind=(mask>0)
spdata= amip_tas
for(t in 1:dim(amip_tas)[3] ){
	ap=amip_tas[,,t]
	ap[!space.ind]=NA
	spdata[,,t]=ap
}
ap= read.amip(spdata[ind1,ind2,],tas_time,lon.out[ind1],lat.out[ind2],6,8)
NE.T=apply(ap,c(3),mean,na.rm=TRUE)
NE.T=mov.detrend(NE.T,len=7)

#(4) Correlation with tas ---------------
pdf(paste('T_tas','.pdf',sep=''),width=6,height=5)
names=c('(a) Jan-Feb-Mar','(b) Feb-Mar-Apr','(c) Mar-Apr-May',	'(d) Apr-May-Jun','(e) May-Jun-Jul','(f) Jun-Jul-Aug')
par(mar=c(1,2,1,3))
par(mfrow=c(2,2))

for (month in c(4,7)){
met= read.amip(amip_tas,tas_time,lon.out,lat.out,month-1,month+1)
for (i in 1:dim(met)[1]){
	for (j in 1:dim(met)[2]){
		met[i,j,]=mov.detrend(met[i,j,],len=7)
	}
}

ap=find.cor2(met, NE.T,p_value=1)
plot.field(ap,lon.out,lat.out,type='sign',mai=c(0.1,0.2,0.2,0.2),legend.mar=4.5,zlim=c(-0.65,0.65))
cr=find.Rlim(0.05,30)
for (i in 1:dim(met)[1]){
	for (j in 1:dim(met)[2]){
		if(abs(ap[i,j])> cr)	points(lon.out[i],lat.out[j],pch=16,cex=0.45,col=1)
	}
}

if(month==4){
	correlation=ap
} else {
	correlation=abind(correlation,ap,along=3)
}
title(names[month-1],cex.main=1)
}
dev.off()
save('correlation','lon.out','lat.out',file='correlation_T-tas.Rdata')

#(4) Correlation with SLP ---------------
pdf(paste('T_slp','.pdf',sep=''),width=6,height=5)
names=c('(a) Jan-Feb-Mar','(b) Feb-Mar-Apr','(c) Mar-Apr-May',	'(d) Apr-May-Jun','(e) May-Jun-Jul','(f) Jun-Jul-Aug')
par(mar=c(1,2,1,3))
par(mfrow=c(2,2))

for (month in c(4,7)){
met= read.amip(amip_psl,slp_time,lon.out,lat.out,month-1,month+1)
for (i in 1:dim(met)[1]){
	for (j in 1:dim(met)[2]){
		met[i,j,]=mov.detrend(met[i,j,],len=7)
	}
}

ap=find.cor2(met, NE.T,p_value=1)
plot.field(ap,lon.out,lat.out,type='sign',mai=c(0.1,0.2,0.2,0.2),legend.mar=4.5,zlim=c(-0.65,0.65))
cr=find.Rlim(0.05,30)
for (i in 1:dim(met)[1]){
	for (j in 1:dim(met)[2]){
		if(abs(ap[i,j])> cr)	points(lon.out[i],lat.out[j],pch=16,cex=0.45,col=1)
	}
}

if(month==4){
	correlation=ap
} else {
	correlation=abind(correlation,ap,along=3)
}
title(names[month-1],cex.main=1)
}
dev.off()
save('correlation','lon.out','lat.out',file='correlation_T-slp.Rdata')

#(4) Correlation with SST ---------------
pdf(paste('T_SST','.pdf',sep=''),width=6,height=5)
names=c('(a) Jan-Feb-Mar','(b) Feb-Mar-Apr','(c) Mar-Apr-May',	'(d) Apr-May-Jun','(e) May-Jun-Jul','(f) Jun-Jul-Aug')
par(mar=c(1,2,1,3))
par(mfrow=c(2,2))

for (month in c(4,7)){
if(month==4) met=MAM_SST
if(month==7) met=JJA_SST
lon.out=lon.frame-360
lat.out=lat.frame
ap=find.cor2(met, NE.T,p_value=1)
plot.field(ap,lon.out,lat.out,type='sign',mai=c(0.1,0.2,0.2,0.2),legend.mar=4.5,zlim=c(-0.65,0.65))
ap[is.na(ap)]=0
cr=find.Rlim(0.05,30)
for (i in 1:dim(met)[1]){
	for (j in 1:dim(met)[2]){
		if(abs(ap[i,j])> cr)	points(lon.out[i],lat.out[j],pch=16,cex=0.45,col=1)
	}
}

if(month==4){
	correlation=ap
} else {
	correlation=abind(correlation,ap,along=3)
}
title(names[month-1],cex.main=1)
}
dev.off()
save('correlation','lon.out','lat.out',file='correlation_T-SST.Rdata')
}